import cv2
import numpy as np
import os
from sam_predict import load_model, do_infer

# 全局变量，用于保存鼠标点击的像素坐标
clicked_point = None

def mouse_callback(event, x, y, flags, param):
    """鼠标点击回调函数"""
    global clicked_point
    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        clicked_point = (x, y)
        print(f"选中的像素点坐标: {clicked_point}")

def get_image_paths(directory):
    # 定義常見圖片格式
    image_extensions = (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".webp")
    image_paths = []

    # 遍歷目錄及其子目錄
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(image_extensions):  # 判斷是否是圖片
                image_paths.append(os.path.join(root, file))
    
    # 對結果進行排序
    image_paths.sort()  # 按完整路徑排序
    return image_paths

def main():
    global clicked_point  # 声明 clicked_point 为全局变量
    checkpoint_path = '/home/JSDC/017254/code/gitee/deep-learing/model/sam_vit_h_4b8939.pth'

    # 使用示例
    directory_path = "/home/JSDC/017254/code/gitee/deep-learing/src_imgs"
    output_txt_path = "/home/JSDC/017254/code/gitee/deep-learing/txt_center/clicked_points.txt"  # 输出的 txt 文件路径

    # 打开文件，准备写入
    with open(output_txt_path, "w") as output_file:
        images = get_image_paths(directory_path)
        # for img in images:
        for index, img in enumerate(images):  # 添加序号
            print(img)
            rgb_image = cv2.imread(img)
            file_name_with_ext = os.path.basename(img)  # 提取文件名（带扩展名）
            file_id = os.path.splitext(file_name_with_ext)[0]  # 去掉扩展名，获取文件 ID
            # 动态生成窗口名称：序号 + 图片名称
            window_name = f"{index + 1}: {file_name_with_ext}"  # 序号从 1 开始
            while True:
                # 显示 RGB 图像
                cv2.namedWindow(window_name)  # 创建窗口
                cv2.setMouseCallback(window_name, mouse_callback)  # 绑定鼠标回调函数
                display_image = rgb_image.copy()

                if clicked_point is not None:
                    try:
                        # 获取点击点的 (x, y)
                        print(f"像素点 {clicked_point}")
                        u, v = clicked_point

                        print(f'file_id: {file_id}')

                        # 写入到文件中
                        output_file.write(f"{file_id}, {u}, {v}\n")
                        output_file.flush()  # 立即将数据写入文件
                        print(f"已写入: {file_id}, {u}, {v}")

                        clicked_point = None  # 重置点击点
                    except ValueError as e:
                        print(e)
                        clicked_point = None  # 重置点击点

                cv2.imshow(window_name, display_image)
                key = cv2.waitKey(1) & 0xFF

                if key == 27:  # 按下 ESC 键退出
                    break

            cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
